import os
import onnx
import copy
import numpy as np
import logging
import onnxruntime
from collections import OrderedDict
from onnx import shape_inference
logging.basicConfig(level=logging.INFO)
from onnx import shape_inference, TensorProto, version_converter, numpy_helper
logger = logging.getLogger("[ONNXOPTIMIZER]")

import math
from matplotlib.pyplot import *
from matplotlib.pyplot import MultipleLocator
decay = 0.9999
step = [x for x in range(299)]
y = [decay * (1 - math.exp(-(x*177) / 2000)) for x in step]
c = 0.19
for a in range(299):
    for b in range(177):
        c = c * 0.9 + 0.1
        print(b, c)
    exit()
# import matplotlib.pyplot as plt
# from matplotlib.pyplot import MultipleLocator
# #从pyplot导入MultipleLocator类，这个类用于设置刻度间隔
#
# plt.plot(step,y,c='green')
# plt.title('Squares',fontsize=24)
# plt.tick_params(axis='both',which='major',labelsize=14)
# plt.xlabel('Numbers',fontsize=14)
# plt.ylabel('Squares',fontsize=14)
# x_major_locator=MultipleLocator(10)
# #把x轴的刻度间隔设置为1，并存在变量里
# y_major_locator=MultipleLocator(0.1)
# #把y轴的刻度间隔设置为10，并存在变量里
# ax=plt.gca()
# #ax为两条坐标轴的实例
# ax.xaxis.set_major_locator(x_major_locator)
# #把x轴的主刻度设置为1的倍数
# ax.yaxis.set_major_locator(y_major_locator)
# #把y轴的主刻度设置为10的倍数
# # plt.xlim(0,10000)
# #把x轴的刻度范围设置为-0.5到11，因为0.5不满一个刻度间隔，所以数字不会显示出来，但是能看到一点空白
# # plt.ylim(-5,110)
# #把y轴的刻度范围设置为-5到110，同理，-5不会标出来，但是能看到一点空白
# plt.show()
exit(0)




def test_model_by_onnxruntime(model):
    logger.info("Test model by onnxruntime")

    input_shape = model.graph.input[0].type.tensor_type.shape.dim

    image_shape = [x.dim_value for x in input_shape]
    image_shape_new = []
    for x in image_shape:
        if x == 0:
            image_shape_new.append(1)
        else:
            image_shape_new.append(x)
    image_shape = image_shape_new
    img_array = np.zeros(image_shape, dtype = np.float32)
    img = img_array
    for node in model.graph.node:
        for output in node.output:
            model.graph.output.extend([onnx.ValueInfoProto(name=output)])
    ort_session = onnxruntime.InferenceSession(model.SerializeToString())
    ort_inputs = {}
    for i, input_ele in enumerate(ort_session.get_inputs()):
        ort_inputs[input_ele.name] = img

    outputs = [x.name for x in ort_session.get_outputs()]
    ort_outs = ort_session.run(outputs, ort_inputs)
    ort_outs = OrderedDict(zip(outputs, ort_outs))
    # logger.info("Test model by onnxruntime success")
    # del model.graph.output[:]
    # model.graph.output.extend(ori_output)
    return ort_outs


if __name__ == "__main__":
    onnx_model_path = "/home/dzhang/work/nbi/testdata/model/mnist.onnx"
    # sess = onnxrt.InferenceSession(onnx_model_path)
    # input_name = sess.get_inputs()[0].name
    # input_shape = sess.get_inputs()[0].shape
    # output_name = sess.get_outputs()[0].name
    # output_shape = sess.get_outputs()[0].shape
    # print(input_name, input_shape)
    # print(output_name, output_shape)
    # x = np.zeros(input_shape, dtype = np.float32)
    # res = sess.run([output_name], {input_name: x})
    # print(res[0].flatten()[:32])
    
    onnx_model = onnx.load(onnx_model_path)
    ort_outs = test_model_by_onnxruntime(onnx_model)
    res = ort_outs["Plus214_Output_0"]
    print(res.flatten()[:32])
    







